from fastapi import FastAPI, HTTPException, Body
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Optional
import numpy as np
import torch
import cv2
import base64
import logging
import traceback
from segment_anything import sam_model_registry, SamPredictor

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# 全局加载模型
try:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    logger.info(f"Using device: {device}")
    
    sam_checkpoint = "/mnt-public/models/meta/sam/sam_vit_h_4b8939.pth"
    model_type = "vit_h"
    
    logger.info("Loading SAM model...")
    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
    sam.to(device=device)
    predictor = SamPredictor(sam)
    logger.info("SAM model loaded successfully")
except Exception as e:
    logger.error(f"Model loading failed: {str(e)}")
    logger.error(traceback.format_exc())
    raise RuntimeError("Failed to load model") from e

class PredictionRequest(BaseModel):
    image: str  # base64编码的图像
    point_coords: Optional[List[List[float]]] = None
    point_labels: Optional[List[int]] = None
    box: Optional[List[float]] = None
    multimask_output: bool = True
    mask_input: Optional[List[List[List[float]]]] = None

def base64_to_image(image_b64: str) -> np.ndarray:
    try:
        image_data = base64.b64decode(image_b64)
        nparr = np.frombuffer(image_data, np.uint8)
        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        if image is None:
            raise ValueError("Failed to decode image")
        return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    except Exception as e:
        logger.error(f"Image decoding error: {str(e)}")
        raise HTTPException(status_code=400, detail="Invalid image data") from e

def validate_prediction_request(request: PredictionRequest):
    """验证预测请求参数的完整性"""
    errors = []
    
    # 检查点坐标和标签的一致性
    if request.point_coords is not None and request.point_labels is None:
        errors.append("point_labels must be provided when point_coords are given")
    if request.point_labels is not None and request.point_coords is None:
        errors.append("point_coords must be provided when point_labels are given")
    
    if request.point_coords is not None and request.point_labels is not None:
        if len(request.point_coords) != len(request.point_labels):
            errors.append("point_coords and point_labels must have the same length")
        
        for i, coord in enumerate(request.point_coords):
            if len(coord) != 2:
                errors.append(f"point_coords[{i}] must have exactly 2 values (x, y)")
    
    # 检查边界框格式
    if request.box is not None:
        if len(request.box) != 4:
            errors.append("box must have exactly 4 values [x1, y1, x2, y2]")
        else:
            x1, y1, x2, y2 = request.box
            if x1 >= x2 or y1 >= y2:
                errors.append("Invalid box coordinates: x2 must be greater than x1, y2 greater than y1")
    
    # 检查掩码输入格式
    if request.mask_input is not None:
        if not isinstance(request.mask_input, list) or not all(isinstance(row, list) for row in request.mask_input):
            errors.append("mask_input must be a 2D list")
    
    # 检查至少有一种输入方式
    if request.point_coords is None and request.box is None and request.mask_input is None:
        errors.append("At least one of point_coords, box or mask_input must be provided")
    
    if errors:
        raise HTTPException(status_code=400, detail="; ".join(errors))

@app.post("/predict")
async def predict(request: PredictionRequest):
    try:
        # 验证请求参数
        validate_prediction_request(request)
        
        # 解码图像
        image_np = base64_to_image(request.image)
        
        # 检查图像尺寸
        if image_np.size == 0:
            raise HTTPException(status_code=400, detail="Empty image after decoding")
        
        # 设置图像
        try:
            predictor.set_image(image_np)
        except Exception as e:
            logger.error(f"Error in set_image: {str(e)}")
            raise HTTPException(status_code=500, detail=f"Image processing error: {str(e)}")
        
        # 准备参数
        params = {"multimask_output": request.multimask_output}
        
        # 处理点坐标和标签
        if request.point_coords and request.point_labels:
            params["point_coords"] = np.array(request.point_coords, dtype=np.float32)
            params["point_labels"] = np.array(request.point_labels, dtype=np.int32)
        
        # 处理边界框
        if request.box:
            # 确保边界框格式正确 [x1, y1, x2, y2]
            if len(request.box) == 4:
                params["box"] = np.array(request.box, dtype=np.float32)[None, :]
            else:
                logger.warning(f"Ignoring invalid box format: {request.box}")
        
        # 处理掩码输入
        if request.mask_input:
            try:
                mask_array = np.array(request.mask_input, dtype=np.float32)
                if mask_array.ndim == 2:
                    mask_array = mask_array[np.newaxis, :, :]
                params["mask_input"] = torch.tensor(mask_array).to(device)
            except Exception as e:
                logger.error(f"Error processing mask_input: {str(e)}")
                # 不抛出异常，继续处理
        
        # 执行预测
        try:
            masks, scores, logits = predictor.predict(**params)
        except Exception as e:
            logger.error(f"Prediction error: {str(e)}")
            logger.error(traceback.format_exc())
            raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")
        
        # 转换结果以便JSON序列化
        masks_list = masks.astype(bool).tolist()
        scores_list = scores.tolist()
        logits_list = logits.tolist() if logits is not None else None
        
        return {
            "masks": masks_list,
            "scores": scores_list,
            "logits": logits_list
        }
    
    except HTTPException as he:
        raise he
    except Exception as e:
        logger.error(f"Unexpected error: {str(e)}")
        logger.error(traceback.format_exc())
        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")

@app.get("/health")
def health_check():
    """健康检查端点"""
    try:
        # 简单检查模型是否已加载
        if predictor is None:
            raise RuntimeError("Predictor not initialized")
        
        # 检查CUDA是否可用（如果是GPU模式）
        if device == "cuda" and not torch.cuda.is_available():
            raise RuntimeError("CUDA not available")
        
        return {"status": "healthy", "device": device}
    except Exception as e:
        logger.error(f"Health check failed: {str(e)}")
        raise HTTPException(status_code=503, detail=f"Service unavailable: {str(e)}")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8001)